import argparse
import os
import yaml
import shutil
import numpy as np
import pandas as pd
from scipy.stats import spearmanr
from tqdm import tqdm

import mindspore as ms
from mindspore import context, nn, Tensor
from mindspore.dataset import GeneratorDataset
from mindspore.train.serialization import load_checkpoint, load_param_into_net
from mindspore_geometric.loader import GraphDataLoader

from src.utils.utils import print_param_num
from model_ms import GNN_model
from data_ms import mutant_dataset

from numpy import nan
from collections import OrderedDict


def mutant_predict(model, loader, protein_names):
    model.set_train(False)
    softmax = nn.Softmax(axis=1)
    protein_num = len(protein_names)
    spear_cor = np.zeros(protein_num)

    for data in tqdm(loader, desc="Mutant Prediction"):
        # Move data to device
        graph_data = data.to_device(ms.context.get_context('DEVICE_ID'))
        out = model(graph_data)
        out = ms.ops.Log(softmax(out[:, :20]) + 1e-9)

        protein_idx = int(data.protein_idx.asnumpy())
        score_info = data.score_info[0]
        num_mutat = len(score_info)

        true_score = Tensor(np.zeros(num_mutat), ms.float32)
        pred_score = Tensor(np.zeros(num_mutat), ms.float32)
        mutat_pt_num = Tensor(np.zeros(num_mutat), ms.int32)

        for mutat_idx in range(num_mutat):
            mutat_info, ts = score_info[mutat_idx]
            true_score[mutat_idx] = ts
            mutat_pt_num[mutat_idx] = len(mutat_info)
            for i in range(len(mutat_info)):
                item = mutat_info[i]
                idx, tgt, src = int(item[1]) - 1, int(item[2]), int(item[0])
                if idx >= out.shape[0]:
                    continue
                pred_score[mutat_idx] += (out[idx, tgt] - out[idx, src])

        df_score = pd.DataFrame({
            "true_score": true_score.asnumpy(),
            "pred_score": pred_score.asnumpy(),
            "mutat_pt_num": mutat_pt_num.asnumpy()
        })
        result_dir = f"result/{args.mutant_dataset.split('/')[-1]}"
        os.makedirs(result_dir, exist_ok=True)
        df_score.to_csv(f"{result_dir}/{protein_names[protein_idx]}.tsv", sep="\t", index=False)

        corr = spearmanr(df_score["true_score"], df_score["pred_score"]).correlation
        spear_cor[protein_idx] = 0 if np.isnan(corr) else corr

    print("-" * 40)
    spear_info = {name: spear_cor[i] for i, name in enumerate(protein_names)}
    print(f"multi_avg: {spear_cor.mean()}")
    return spear_info, spear_cor.mean()


def create_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--use_sasa", action="store_true")
    parser.add_argument("--use_bfactor", action="store_true")
    parser.add_argument("--use_dihedral", action="store_true")
    parser.add_argument("--use_coordinate", action="store_true")
    parser.add_argument("--gnn", type=str, default="egnn")
    parser.add_argument("--problem_type", type=str, default="aa_classification")
    parser.add_argument("--num_classes", type=int, default=20)
    parser.add_argument("--node_dim", type=int, default=26)
    parser.add_argument("--edge_dim", type=int, default=93)
    parser.add_argument("--layer_num", type=int, default=6)
    parser.add_argument("--dropout", type=float, default=0.0)
    parser.add_argument("--checkpoint", type=str, default="")
    parser.add_argument("--c_alpha_max_neighbors", type=int, default=10)
    parser.add_argument("--protein_dataset", type=str, default="data/cath40_k10_dyn_imem")
    parser.add_argument("--mutant_dataset", type=str, default="data/evaluation")
    parser.add_argument("--gnn_config", type=str, default="src/Egnnconfig/egnn.yaml")
    return parser.parse_args()


if __name__ == '__main__':
    args = create_parser()

    # Setup context
    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")

    # Load config
    with open(args.gnn_config, 'r') as f:
        cfg = yaml.load(f, Loader=yaml.FullLoader)
    gnn_config = cfg[args.gnn]

    # Load dataset
    mutant = mutant_dataset(args)
    protein_names = mutant.get_names()
    loader = GraphDataLoader(mutant, batch_size=1, shuffle=False)

    # Build model
    model = GNN_model(gnn_config, args)
    print_param_num(model)

    # Load checkpoint
    param_dict = load_checkpoint(args.checkpoint)
    new_dict = {f"GNN_model.{k}": v for k, v in param_dict.items()}
    fixed_dict = OrderedDict()
    for k, v in new_dict.items():
        if k.startswith('GNN_model.GNN_model.'):
            fixed_key = k[len('GNN_model.'):]
        else:
            fixed_key = k
        fixed_dict[fixed_key] = v
    load_param_into_net(model, fixed_dict)

    # Run prediction
    spear_info, multi_mean = mutant_predict(model, loader, protein_names)
